import numpy as np
import matplotlib.pyplot as plt


class KalmanFilter:
    def __init__(self, dt, state_dim, measure_dim):
        """
        初始化卡尔曼滤波器参数
        - dt: 时间步长
        - state_dim: 状态向量维度（例如：[位置, 速度]）
        - measure_dim: 观测向量维度（例如：直接测量的位置）
        """
        self.dt = dt

        # 状态转移矩阵（假设匀速模型）
        self.F = np.array([[1, dt],
                           [0, 1]])

        # 观测矩阵（仅观测位置）
        self.H = np.array([[1, 0]])

        # 过程噪声协方差矩阵（调整系统不确定性）
        self.Q = np.eye(state_dim) * 0.01

        # 观测噪声协方差矩阵（调整传感器误差）
        self.R = np.eye(measure_dim) * 1

        # 初始状态协方差矩阵
        self.P = np.eye(state_dim)

        # 初始状态向量（例如：[位置, 速度]）
        self.x = np.zeros((state_dim, 1))

    def predict(self):
        # 预测状态和协方差
        self.x = self.F @ self.x
        self.P = self.F @ self.P @ self.F.T + self.Q
        return self.x

    def update(self, z):
        # 计算卡尔曼增益
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)

        # 更新状态和协方差
        self.x = self.x + K @ (z - self.H @ self.x)
        self.P = (np.eye(self.P.shape[0]) - K @ self.H) @ self.P
        return self.x


# 生成模拟数据
np.random.seed(42)
n_steps = 50
true_position = np.linspace(0, 10, n_steps)  # 真实位置（匀速运动）
measurements = true_position + np.random.normal(0, 1, n_steps)  # 带噪声的观测值

# 初始化卡尔曼滤波器
kf = KalmanFilter(dt=0.1, state_dim=2, measure_dim=1)
estimated_states = []

# 迭代滤波
for z in measurements:
    # 预测
    kf.predict()
    # 更新
    estimated_state = kf.update(z.reshape(-1, 1))
    estimated_states.append(estimated_state[0, 0])

# 可视化结果
plt.figure(figsize=(10, 6))
plt.plot(true_position, label='True Position', linestyle='--')
plt.plot(measurements, label='Measurements', marker='o', alpha=0.5)
plt.plot(estimated_states, label='Kalman Filter', marker='x', color='red')
plt.xlabel('Time Step')
plt.ylabel('Position')
plt.legend()
plt.title('Kalman Filter for Position Tracking')
plt.show()